import os
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import random

import imageio
from torch import nn

from collections import defaultdict
import argparse
import cv2
import numpy as np
import torch

from pytorch_grad_cam import GradCAM, \
    ScoreCAM, \
    GradCAMPlusPlus, \
    AblationCAM, \
    XGradCAM, \
    EigenCAM, \
    EigenGradCAM, \
    LayerCAM, \
    FullGrad

from pytorch_grad_cam.utils.image import show_cam_on_image, \
    preprocess_image
from pytorch_grad_cam.ablation_layer import AblationLayerVit
from model import swin_base_patch4_window12_384_in22k as Model
from config_hyperparam import cfg

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--use-cuda', action='store_true', default=True)
    parser.add_argument(
        '--image-path',
        type=str,
        default=cfg.dataset_path,
        help='Input image path')
    parser.add_argument('--aug_smooth', action='store_true')
    parser.add_argument(
        '--eigen_smooth',
        action='store_true')

    parser.add_argument(
        '--method',
        type=str,
        default='gradcam',)

    args = parser.parse_args()
    args.use_cuda = args.use_cuda and torch.cuda.is_available()
    if args.use_cuda:
        print('Using GPU for acceleration')
    else:
        print('Using CPU for computation')

    return args


def reshape_transform(tensor, height=7, width=7):
    print(tensor.shape)
    result = tensor.reshape(tensor.size(0),
                            height, width, tensor.size(2))

    result = result.transpose(2, 3).transpose(1, 2)
    return result


def get_dataset(root_dir):
    class_to_samples = defaultdict(list)
    for cls_name in os.listdir(root_dir):  
        cls_dir = os.path.join(root_dir, cls_name)  
        if os.path.isdir(cls_dir):  
            for f in os.listdir(cls_dir):  
                if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):  
                    path = os.path.join(cls_dir, f)  
                    class_to_samples[cls_name].append(path)  
    random.seed(123)
    samples = []  
    for cls_name, s in class_to_samples.items():  
        random.shuffle(s)  
        split_index = int(len(s) * 0.8)  
        samples.extend([path for path in s[split_index:]])
    #print(samples)
    return samples

def gen_heatmap(test_img_path, cam):
    for i, tip in enumerate(test_img_path):
        img = cv2.imread(tip, 1)  
        #if img is None:
            #continue
        rgb_img = cv2.imread(tip, 1)[:, :, ::-1]
        rgb_img = cv2.resize(rgb_img, (224, 224))
        rgb_img = np.float32(rgb_img) / 255
        input_tensor = preprocess_image(rgb_img, mean=[0.5, 0.5, 0.5],
                                        std=[0.5, 0.5, 0.5])
        #print(input_tensor.shape)
        cam.batch_size = 32

        grayscale_cam = cam(input_tensor=input_tensor,
                            targets=None,
                            eigen_smooth=args.eigen_smooth,
                            aug_smooth=args.aug_smooth)

        grayscale_cam = grayscale_cam[0, :]

        cam_image = show_cam_on_image(rgb_img, grayscale_cam)
        cv2.imwrite(f'./checkpoint-catDog/heatmap/cam_{i}.jpg', cam_image)
        print(f'cam_{i}.jpg save')

if __name__ == '__main__':


    args = get_args()
    methods = \
        {"gradcam": GradCAM,
         "scorecam": ScoreCAM,
         "gradcam++": GradCAMPlusPlus,
         "ablationcam": AblationCAM,
         "xgradcam": XGradCAM,
         "eigencam": EigenCAM,
         "eigengradcam": EigenGradCAM,
         "layercam": LayerCAM,
         "fullgrad": FullGrad}

    if args.method not in list(methods.keys()):
        raise Exception(f"method should be one of {list(methods.keys())}")

    #model = Model(cfg.num_class)
    #model = nn.DataParallel(model, device_ids=cfg.gpu_idx)
    model = Model(cfg.num_class).to(cfg.device)
    checkpoint = torch.load(cfg.best_model_path)  # 获取日志
    print(model.load_state_dict(checkpoint["model"]))

    model.eval()

    if args.use_cuda:
        model = model.cuda()

    # target_layers = [model.module.layers[-1].blocks[-1].norm2]
    #target_layers = [model.module.selector]
    target_layers = [model.selector]

    if args.method not in methods:
        raise Exception(f"Method {args.method} not implemented")

    if args.method == "ablationcam":
        cam = methods[args.method](model=model,
                                   target_layers=target_layers,
                                   use_cuda=args.use_cuda,
                                   reshape_transform=reshape_transform,
                                   ablation_layer=AblationLayerVit())
    else:
        cam = methods[args.method](model=model,
                                   target_layers=target_layers,

                                   reshape_transform=reshape_transform)

    test_img_path = get_dataset(args.image_path) # 图片路径
    test_img_path = random.sample(test_img_path, 500)  # 随机抽n个画图
    gen_heatmap(test_img_path, cam)
